Count square submatrices with all ones

Time: O(MxN); Space: O(1); medium

Given a m * n matrix of ones and zeros, return how many square submatrices have all ones.

Example 1:

Input: matrix =

[
  [0,1,1,1],
  [1,1,1,1],
  [0,1,1,1]
]

Output: 15

Explanation:

  • There are 10 squares of side 1.

  • There are 4 squares of side 2.

  • There is 1 square of side 3.

  • Total number of squares = 10 + 4 + 1 = 15.

Example 2:

Input: matrix =

[
  [1,0,1],
  [1,1,0],
  [1,1,0]
]

Output: 7

Explanation:

  • There are 6 squares of side 1.

  • There is 1 square of side 2.

  • Total number of squares = 6 + 1 = 7.

Constraints:

  • 1 <= len(arr) <= 300

  • 1 <= len(arr[0]) <= 300

  • 0 <= arr[i][j] <= 1

Hints:

  1. Create an additive table that counts the sum of elements of submatrix with the superior corner at (0,0).

  2. Loop over all subsquares in O(n^3) and check if the sum make the whole array to be ones, if it checks then add 1 to the answer.

1. Dynamic programming [O(MxN), O(1)]

[4]:

class Solution1(object):
    """
    Time: O(M*N)
    Space: O(1)
    """
    def countSquares(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: int
        """
        for i in range(1, len(matrix)):
            for j in range(1, len(matrix[0])):
                if not matrix[i][j]:
                    continue
                l = min(matrix[i-1][j], matrix[i][j-1])
                matrix[i][j] = l+1 if matrix[i-l][j-l] else l

        return sum(x for row in matrix for x in row)

[5]:

s = Solution1()

matrix = [
  [0,1,1,1],
  [1,1,1,1],
  [0,1,1,1]
]
assert s.countSquares(matrix) == 15

matrix = [
  [1,0,1],
  [1,1,0],
  [1,1,0]
]
assert s.countSquares(matrix) == 7